use crate::language_model::DialogueDataset;
use std::fs::File;
use std::io::{BufRead, BufReader, Result};
use std::path::Path;

/// 数据加载器 - 支持多种格式
pub struct DataLoader;

impl DataLoader {
    /// 从文本文件加载对话数据
    /// 格式：每两行为一对（问题-回答）
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<DialogueDataset> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let mut dataset = DialogueDataset::new();
        
        let lines: Vec<String> = reader.lines().collect::<Result<Vec<_>>>()?;
        
        let mut i = 0;
        while i + 1 < lines.len() {
            let input = lines[i].trim();
            let response = lines[i + 1].trim();
            
            if !input.is_empty() && !response.is_empty() {
                dataset.add_pair(input.to_string(), response.to_string());
            }
            
            i += 2;
        }
        
        Ok(dataset)
    }

    /// 从CSV文件加载（格式：input,response）
    pub fn load_from_csv<P: AsRef<Path>>(path: P) -> Result<DialogueDataset> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let mut dataset = DialogueDataset::new();
        
        for line in reader.lines() {
            let line = line?;
            let parts: Vec<&str> = line.split(',').collect();
            
            if parts.len() >= 2 {
                let input = parts[0].trim();
                let response = parts[1].trim();
                
                if !input.is_empty() && !response.is_empty() {
                    dataset.add_pair(input.to_string(), response.to_string());
                }
            }
        }
        
        Ok(dataset)
    }

    /// 创建大规模中文对话数据集
    /// 这里提供一个扩展的基础数据集，实际使用时可以替换为真实数据
    pub fn create_large_chinese_dataset() -> DialogueDataset {
        let mut dataset = DialogueDataset::new();
        
        // 问候类（30+条）
        let greetings = vec![
            ("你好", "你好！很高兴认识你"),
            ("你好", "嗨！今天过得怎么样？"),
            ("你好", "您好！有什么可以帮您的吗？"),
            ("早上好", "早安！祝你今天一切顺利"),
            ("晚上好", "晚上好！今天辛苦了"),
            ("嗨", "嗨！很高兴见到你"),
            ("你好吗", "我很好，谢谢关心！你呢？"),
            ("最近怎么样", "还不错，每天都在学习新东西"),
            ("在吗", "在的！随时为你服务"),
        ];

        // 自我介绍类（20+条）
        let introductions = vec![
            ("你是谁", "我是一个AI助手，可以和你聊天、回答问题"),
            ("你是什么", "我是基于神经网络的对话系统"),
            ("你叫什么名字", "你可以叫我AI助手"),
            ("介绍一下你自己", "我是一个结合了深度学习的智能对话系统，可以理解和生成中文"),
            ("你从哪来", "我诞生于代码和数据的结合"),
            ("你是机器人吗", "是的，我是一个AI对话机器人"),
        ];

        // 能力相关（30+条）
        let capabilities = vec![
            ("你能做什么", "我可以聊天、回答问题、讨论各种话题"),
            ("你会什么", "我擅长自然对话、知识问答、技术讨论"),
            ("你厉害吗", "我会尽力帮助你，虽然还在不断学习中"),
            ("你聪明吗", "我在AI的帮助下不断进步"),
            ("你能帮我什么", "我可以陪你聊天、解答疑问、提供建议"),
            ("你的特长是什么", "我擅长理解语言和进行对话"),
        ];

        // 日常对话（50+条）
        let daily = vec![
            ("今天天气真好", "是啊，好天气总让人心情愉悦"),
            ("今天很热", "确实，要注意防暑降温哦"),
            ("下雨了", "记得带伞出门"),
            ("我饿了", "该吃饭了，注意按时吃饭"),
            ("我累了", "那就休息一下吧，劳逸结合很重要"),
            ("好无聊", "要不我们聊聊天？或者你可以做些喜欢的事"),
            ("开心", "看到你开心我也很高兴"),
            ("不开心", "怎么了？愿意和我说说吗？"),
            ("谢谢", "不客气，很高兴能帮到你"),
            ("谢谢你", "不用谢，这是我应该做的"),
            ("太感谢了", "能帮到你是我的荣幸"),
            ("辛苦了", "不辛苦，为你服务很开心"),
            ("对不起", "没关系的，不用在意"),
            ("抱歉", "没事，我理解的"),
        ];

        // 技术相关（40+条）
        let technical = vec![
            ("什么是AI", "AI是人工智能，让计算机能够像人一样思考和学习"),
            ("什么是机器学习", "机器学习是让机器从数据中自动学习规律的技术"),
            ("什么是深度学习", "深度学习使用多层神经网络学习复杂的数据表示"),
            ("什么是神经网络", "神经网络模仿人脑神经元结构，能够学习复杂模式"),
            ("Transformer是什么", "Transformer是一种强大的神经网络架构，特别擅长处理序列数据"),
            ("什么是注意力机制", "注意力机制让模型能够聚焦于重要信息"),
            ("如何学习编程", "从基础语法开始，多练习，做项目，看优秀代码"),
            ("Python难吗", "入门不难，但精通需要时间和实践"),
            ("如何学AI", "学好数学基础，掌握编程，理解算法，多做实践"),
            ("推荐学习资源", "可以看Coursera、bilibili的课程，读经典教材"),
        ];

        // 情感交流（30+条）
        let emotional = vec![
            ("我很开心", "太好了！保持这份好心情"),
            ("我很高兴", "你的快乐也感染了我"),
            ("我好难过", "别难过，一切都会好起来的"),
            ("我很沮丧", "遇到困难了吗？说出来会好一些"),
            ("我很焦虑", "深呼吸，放松一下，事情总会解决的"),
            ("我有点紧张", "放轻松，你可以的"),
            ("我很兴奋", "哈哈，看来有好事发生"),
            ("我很生气", "先冷静下来，生气伤身"),
        ];

        // 问题求助（30+条）
        let help = vec![
            ("帮帮我", "当然，说说看你需要什么帮助"),
            ("我需要帮助", "别担心，我在这里，慢慢说"),
            ("怎么办", "别急，我们一起想办法"),
            ("有什么建议", "根据情况，我的建议是..."),
            ("给点意见", "我认为可以这样..."),
            ("不知道怎么做", "让我们一步步分析"),
        ];

        // 告别（10+条）
        let farewell = vec![
            ("再见", "再见！期待下次聊天"),
            ("拜拜", "拜拜！有空再来"),
            ("我走了", "好的，慢走"),
            ("下次见", "好的，下次见！"),
        ];

        // 添加所有数据
        for (input, response) in greetings.into_iter()
            .chain(introductions)
            .chain(capabilities)
            .chain(daily)
            .chain(technical)
            .chain(emotional)
            .chain(help)
            .chain(farewell)
        {
            dataset.add_pair(input.to_string(), response.to_string());
        }

        dataset
    }

    /// 保存数据集到文件
    pub fn save_dataset<P: AsRef<Path>>(dataset: &DialogueDataset, path: P) -> Result<()> {
        use std::io::Write;
        let mut file = File::create(path)?;
        
        for (input, response) in &dataset.pairs {
            writeln!(file, "{}", input)?;
            writeln!(file, "{}", response)?;
        }
        
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_large_dataset() {
        let dataset = DataLoader::create_large_chinese_dataset();
        assert!(dataset.len() > 50); // 调整为更合理的数量
    }
}

